import os
from dotenv import load_dotenv

from langchain_openai import ChatOpenAI
from langchain_core.tools import tool
from langgraph.prebuilt import create_react_agent

from my_deep_agents_from_scratch.state import DeepAgentState
from my_deep_agents_from_scratch.prompts import SUBAGENT_USAGE_INSTRUCTIONS
from my_deep_agents_from_scratch.task_tool import _create_task_tool
from utils import format_messages, show_prompt

from datetime import datetime

def main():
    load_dotenv(os.path.join("..", ".env"), override=True)
    model = os.getenv("MODEL")
    base_url = os.getenv("BASE_URL")
    api_key = os.getenv("_API_KEY")
    # show_prompt(SUBAGENT_USAGE_INSTRUCTIONS)

    # Limits
    max_concurrent_research_units = 3
    max_researcher_iterations = 3

    # Mock search result
    search_result = """模型上下文协议（MCP）是由Anthropic开发的一种标准协议，旨在实现AI模型与外部系统（如工具、数据库及其他服务）的无缝集成。
    它作为标准化的通信层，使AI模型能够以一致高效的方式访问并利用来自不同数据源的信息。
    本质上，MCP通过为数据交换提供统一语言，极大简化了将AI助手连接到外部服务的过程。"""

    # Mock search tool
    @tool(parse_docstring=True)
    def web_search(
            query: str,
    ):
        """Search the web for information on a specific topic.

        This tool performs web searches and returns relevant results
        for the given query. Use this when you need to gather information from
        the internet about any topic.

        Args:
            query: The search query string. Be specific and clear about what
                   information you're looking for.

        Returns:
            Search results from the search engine.

        Example:
            web_search("machine learning applications in healthcare")
        """
        return search_result

    # Add mock research instructions
    SIMPLE_RESEARCH_INSTRUCTIONS = """You are a researcher. Research the topic provided to you. IMPORTANT: Just make a single call to the web_search tool and use the result provided by the tool to answer the provided topic."""

    # Create research sub-agent
    research_sub_agent = {
        "name": "research-agent",
        "description": "Delegate research to the sub-agent researcher. Only give this researcher one topic at a time.",
        "prompt": SIMPLE_RESEARCH_INSTRUCTIONS,
        "tools": ["web_search"],
    }

    # Create agent using create_react_agent directly
    # model = ChatOpenAI(model=model, base_url=base_url, api_key=api_key)
    from langchain_deepseek.chat_models import ChatDeepSeek
    model = ChatDeepSeek(model=model, temperature=0, api_key=api_key, base_url=base_url)

    # Tools for sub-agent
    sub_agent_tools = [web_search]

    # Create task tool to delegate tasks to sub-agents
    task_tool = _create_task_tool(
        sub_agent_tools, [research_sub_agent], model, DeepAgentState
    )

    # Tools
    delegation_tools = [task_tool]

    # Create agent with system prompt
    agent = create_react_agent(
        model,
        delegation_tools,
        prompt=SUBAGENT_USAGE_INSTRUCTIONS.format(
            max_concurrent_research_units=max_concurrent_research_units,
            max_researcher_iterations=max_researcher_iterations,
            date=datetime.now().strftime("%a %b %d, %Y"),
        ),
        state_schema=DeepAgentState,
    )

    # # Show the agent
    # png_data = agent.get_graph(xray=True).draw_mermaid_png()
    # output_path = "3_subagents_graph.png"  # 保存到文件
    # with open(output_path, "wb") as f:
    #     f.write(png_data)

    result = agent.invoke(
        {
            "messages": [
                {
                    "role": "user",
                    "content": "给我一份模型上下文协议（MCP）的简要概述。",
                }
            ],
        }
    )

    format_messages(result["messages"])

if __name__ == '__main__':
    main()